# This file defines classes to hold results from a RAG system which need to be evaluated.
from dataclasses import dataclass

@dataclass
class RetrievalResult:
    """This class holds the output from a retrieval system."""
    # The query that was used to retrieve the passages
    query: str
    # The passages that were retrieved, each entry in the dict is a passage
    # with the key being the passage ID.
    retrieved_passages: dict[str, str]

@dataclass
class GeneratedAnswerPart:
    # A piece of cited text generated by the RAG system.
    text: str
    # Passage citations for the text.
    citations: list[str]

@dataclass
class AugmentedGenerationResult:
    """This class holds the output from a generation system."""
    # The query that was used to generate the answer
    query: str
    # The generated answer along with citations.
    generated_answer: list[GeneratedAnswerPart]

@dataclass
class RAGResult:
    """This class holds the output from a RAG system."""
    retrieval_result: RetrievalResult
    generation_result: AugmentedGenerationResult
